﻿using LLama.Common;
using LLama;
using System.Text;

namespace LLamaSharpTest;

public class Chat
{
    ChatSession session;




    InferenceParams inferenceParams = new InferenceParams()
    {
        Temperature = 0.9f,
        AntiPrompts = new List<string> { "用户：" }
    };


    private string ConvertEncoding(string input, Encoding original, Encoding target)
    {
        byte[] bytes = original.GetBytes(input);
        var convertedBytes = Encoding.Convert(original, target, bytes);
        return target.GetString(convertedBytes);
    }


    public void Init()
    {
        // Register provider for GB2312 encoding
        Encoding.RegisterProvider(CodePagesEncodingProvider.Instance);


        Console.ForegroundColor = ConsoleColor.Yellow;
        Console.WriteLine("This example shows how to use Chinese with gb2312 encoding, which is common in windows. It's recommended" +
            " to use https://huggingface.co/hfl/chinese-alpaca-2-7b-gguf/blob/main/ggml-model-q5_0.gguf, which has been verified by LLamaSharp developers.");
        Console.ForegroundColor = ConsoleColor.White;


        string modelPath = @"E:\LMModels\ay\Repository\qwen1_5-7b-chat-q8_0.gguf";// @"E:\Models\llama-2-7b-chat.Q4_K_M.gguf";


        var parameters = new ModelParams(modelPath)
        {
            ContextSize = 1024,
            Seed = 1337,
            //GpuLayerCount = 5,
            Encoding = Encoding.UTF8
        };


        var model = LLamaWeights.LoadFromFile(parameters);
        var context = model.CreateContext(parameters);
        var executor = new InteractiveExecutor(context);
        var chatHistory = new ChatHistory();
        session = new ChatSession(executor, chatHistory);
        session
            .WithHistoryTransform(new LLamaTransforms.DefaultHistoryTransform());
    }


    public async Task Run(string userInput, Action<string> callback)
    {
        while (userInput != "exit")
        {
            userInput = ConvertEncoding(userInput, Encoding.GetEncoding("gb2312"), Encoding.UTF8);


            if (userInput == "save")
            {
                session.SaveSession("chat-with-kunkun-chinese");


            }
            else if (userInput == "regenerate")
            {
                await foreach (
                    var text
                    in session.RegenerateAssistantMessageAsync(
                        inferenceParams))
                {
                    callback(ConvertEncoding(text, Encoding.UTF8, Encoding.GetEncoding("gb2312")));
                }
            }
            else
            {
                await foreach (
                    var text
                    in session.ChatAsync(
                        new ChatHistory.Message(AuthorRole.User, userInput),
                        inferenceParams))
                {
                    callback(text);
                }
            }


            userInput = "";
        }
    }
}